{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UYyFIEKEkmHb"
   },
   "source": [
    "# Dria\n",
    "\n",
    ">[Dria](https://dria.co/) is a hub of public RAG models for developers to both contribute and utilize a shared embedding lake. This notebook demonstrates how to use the `Dria API` for data retrieval tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VNTFUgK9kmHd"
   },
   "source": [
    "# Installation\n",
    "\n",
    "Ensure you have the `dria` package installed. You can install it using pip:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X--1A8EEkmHd"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet dria"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xRbRL0SgkmHe"
   },
   "source": [
    "# Configure API Key\n",
    "\n",
    "Set up your Dria API key for access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "hGqOByNMkmHe"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"DRIA_API_KEY\"] = \"DRIA_API_KEY\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nDfAEqQtkmHe"
   },
   "source": [
    "# Initialize Dria Retriever\n",
    "\n",
    "Create an instance of `DriaRetriever`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vlyorgCckmHe"
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers import DriaRetriever\n",
    "\n",
    "api_key = os.getenv(\"DRIA_API_KEY\")\n",
    "retriever = DriaRetriever(api_key=api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j7WUY5jBOLQd"
   },
   "source": [
    "# **Create Knowledge Base**\n",
    "\n",
    "Create a knowledge on [Dria's Knowledge Hub](https://dria.co/knowledge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L5ER81eWOKnt"
   },
   "outputs": [],
   "source": [
    "contract_id = retriever.create_knowledge_base(\n",
    "    name=\"France's AI Development\",\n",
    "    embedding=DriaRetriever.models.jina_embeddings_v2_base_en.value,\n",
    "    category=\"Artificial Intelligence\",\n",
    "    description=\"Explore the growth and contributions of France in the field of Artificial Intelligence.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9VCTzSFpkmHe"
   },
   "source": [
    "# Add Data\n",
    "\n",
    "Load data into your Dria knowledge base."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xeTMafIekmHf"
   },
   "outputs": [],
   "source": [
    "texts = [\n",
    "    \"The first text to add to Dria.\",\n",
    "    \"Another piece of information to store.\",\n",
    "    \"More data to include in the Dria knowledge base.\",\n",
    "]\n",
    "\n",
    "ids = retriever.add_texts(texts)\n",
    "print(\"Data added with IDs:\", ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dy1UlvLCkmHf"
   },
   "source": [
    "# Retrieve Data\n",
    "\n",
    "Use the retriever to find relevant documents given a query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9y3msv9tkmHf"
   },
   "outputs": [],
   "source": [
    "query = \"Find information about Dria.\"\n",
    "result = retriever.invoke(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
